from __future__ import print_function
#%matplotlib inline
import configargparse
import os
import random
import torch
import os.path as osp

import torch.utils.data
from torch.utils.data import Dataset, DataLoader

import numpy as np

import cv2

import warnings
from torch.serialization import SourceChangeWarning

warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=SourceChangeWarning)

from numba.core.errors import NumbaPerformanceWarning
import warnings

warnings.simplefilter('ignore', category=NumbaPerformanceWarning)

# Import inference module
from infer_get_seq import my_inference

if __name__ == "__main__":

    # Model currently developed to work on GPU
    device = torch.device("cuda:0" if (torch.cuda.is_available()) else "cpu")
    print("* Working on " + str(device))
    assert str(device) == "cuda:0"

    # Load model_generator
    generator = torch.load('models/model_generator_217.pth')

    # Basic inference

    save_dir = 'odisr_seq_x4'
    hr_dir = '/home/mwh/mwh/dataset/lau_dataset/odisr/training/HR'
    lr_dir = '/home/mwh/mwh/dataset/lau_dataset/odisr/training/LR/X4'
    scale = 4

    lr_list = [osp.join(lr_dir, name) for name in sorted(os.listdir(lr_dir))]
    hr_list = [osp.join(hr_dir, name) for name in sorted(os.listdir(hr_dir))]

    for hr_path, lr_path in zip(hr_list, lr_list):
        hr_img = cv2.imread(hr_path)
        lr_img = cv2.imread(lr_path)

        lr_h = lr_img.shape[0] // 2
        lr_w = lr_img.shape[1] // 2
        lr_step = 32

        hr_h = lr_h * scale
        hr_w = lr_w * scale

        lr_sub_dir = osp.join(save_dir, 'LR',
                              lr_path.split('/')[-1].split('.')[0])
        os.makedirs(lr_sub_dir, exist_ok=True)

        hr_sub_dir = osp.join(save_dir, 'HR',
                              hr_path.split('/')[-1].split('.')[0])
        os.makedirs(hr_sub_dir, exist_ok=True)

        idx = 0
        point_scan = my_inference(image_path=lr_path,
                                  generator=generator,
                                  device=device,
                                  n_generated=4)
        lr_offset_h_list = [
            i for i in range(0, lr_img.shape[0] - lr_h + 1, lr_step)
        ]
        lr_offset_w_list = [i for i in range(0, lr_img.shape[1] + 1, lr_step)]
        point_selected = []
        for lr_offset_h in lr_offset_h_list:
            for lr_offset_w in lr_offset_w_list:
                point_selected.append([lr_offset_h, lr_offset_w])
        point_selected.extend(point_scan)
        point_selected = sorted(point_selected)

        for lr_offset_h, lr_offset_w in point_selected:

            # 截取 lr_part
            lr_part = np.zeros([lr_h, lr_w, 3])
            for i in range(lr_offset_h, lr_offset_h + lr_h):
                for j in range(lr_offset_w, lr_offset_w + lr_w):
                    lr_part[i - lr_offset_h,
                            j - lr_offset_w, :] = lr_img[i % lr_img.shape[0],
                                                         j % lr_img.shape[1], :]
            # 截取 hr_part
            hr_part = np.zeros([hr_h, hr_w, 3])
            for i in range(lr_offset_h * scale, lr_offset_h * scale + hr_h):
                for j in range(lr_offset_w * scale, lr_offset_w * scale + hr_w):
                    hr_part[i - lr_offset_h * scale, j - lr_offset_w *
                            scale, :] = hr_img[i % (lr_img.shape[0] * scale),
                                               j % (lr_img.shape[1] * scale), :]

            img_name = f'{idx:03}.png'
            cv2.imwrite(osp.join(lr_sub_dir, img_name), lr_part)
            cv2.imwrite(osp.join(hr_sub_dir, img_name), hr_part)
            idx += 1
        print(f"{lr_path} : {hr_path} Done.")
    print("Done.")
